﻿using System;
using System.Collections.Generic;
using System.Windows.Automation;
using QuickUIT.Utilities;

namespace QuickUIT
{
    /// <summary>
    /// Provides functionality for finding controls in the automation tree. 
    /// </summary>
    public static class AutomationElementFinder
    {
        /// <summary>
        /// Gets a window control. 
        /// </summary>
        /// <param name="property">The property to search by.</param>
        /// <param name="value">The value of the property to match.</param>
        /// <returns>An <see cref="AutomationElement"/> for a window. If no window is found, a null reference is returned.</returns>
        public static AutomationElement GetWindow(AutomationProperty property, object value)
        {
            AutomationElement foundWindow = null;
            var searchCriteria = new AutomationElementSearchCriteria();
            searchCriteria.Add(new AutomationElementExactSearchCondition(AutomationElementIdentifiers.ControlTypeProperty, ControlType.Window));
            searchCriteria.Add(new AutomationElementExactSearchCondition(property, value));
            FindFirst(AutomationElement.RootElement, searchCriteria);
            return foundWindow;
        }

        /// <summary>
        /// Retrieves the parent window that contains the specified UI Automation element.
        /// </summary>
        /// <param name="element">The element whose parent window is to be returned.</param>
        /// <returns>The parent window, or a null reference if the parent window cannot be found.</returns>
        public static AutomationElement GetParentWindow(AutomationElement element)
        {
            AutomationElement elementParent = element;
            do
            {
                elementParent = TreeWalker.ControlViewWalker.GetParent(elementParent);
                if (elementParent != null && elementParent.Current.ControlType == ControlType.Window)
                {
                    return elementParent;
                }
            }
            while (elementParent != null);
            return elementParent;
        }

        /// <summary>
        /// Returns the first element that matches the specified search conditions.
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="searchCriteria">The object containing the criteria to match.</param>
        /// <returns>The first element that satisfies the conditions, or Nothing if no match is found.</returns>
        public static AutomationElement FindFirst(AutomationElement parentElement, IAutomationElementSearchCriteria searchCriteria)
        {
            return FindFirst(parentElement, searchCriteria, true);
        }

        /// <summary>
        /// Returns the first element that matches the specified search conditions.
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="searchCriteria">The object containing the criteria to match.</param>
        /// <param name="excludeParent">Exclude the parent element from the search.</param>
        /// <returns>The first element that satisfies the conditions, or Nothing if no match is found.</returns>
        public static AutomationElement FindFirst(AutomationElement parentElement, IAutomationElementSearchCriteria searchCriteria, bool excludeParent)
        {
            var foundElements = Find(parentElement, searchCriteria, true, excludeParent);
            return foundElements.Count == 0 ? null : foundElements[0];
        }

        /// <summary>
        /// Returns the first element that matches the specified search conditions within a specified allowable timeframe. 
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="searchCriteria">The object containing the criteria to match.</param>
        /// <param name="milliseconds">The maximum time in milliseconds to allow for attempting to find an element.</param>
        /// <returns>The first element that satisfies the conditions, or Nothing if no match is found within the alloted time.</returns>
        public static AutomationElement FindFirst(AutomationElement parentElement, IAutomationElementSearchCriteria searchCriteria, int milliseconds)
        {
            return FindFirst(parentElement, searchCriteria, milliseconds, true); 
        }

        /// <summary>
        /// Returns the first element that matches the specified search conditions within a specified allowable timeframe. 
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="searchCriteria">The object containing the criteria to match.</param>
        /// <param name="milliseconds">The maximum time in milliseconds to allow for attempting to find an element.</param>
        /// <param name="excludeParent">Exclude the parent element from the search.</param>
        /// <returns>The first element that satisfies the conditions, or Nothing if no match is found within the alloted time.</returns>
        public static AutomationElement FindFirst(AutomationElement parentElement, IAutomationElementSearchCriteria searchCriteria, int milliseconds, bool excludeParent)
        {
            var timer = new AutomationTimer(milliseconds);
            timer.Start();
            while (!timer.HasTimedOut())
            {
                var foundElements = Find(parentElement, searchCriteria, true, excludeParent);
                if (foundElements.Count > 0)
                {
                    return foundElements[0];
                }
            }
            return null;
        }

        /// <summary>
        /// Returns all descendant elements that match the specified search conditions.  
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="searchCriteria">The object containing the criteria to match.</param>
        /// <returns>A collection of elements that satisfy the specified conditions. 
        /// If there are no matches, an empty collection is returned.</returns>
        public static List<AutomationElement> FindAll(AutomationElement parentElement, IAutomationElementSearchCriteria searchCriteria)
        {
            return Find(parentElement, searchCriteria, false, true);
        }

        /// <summary>
        /// Returns all descendant elements that match the specified search conditions.  
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="searchCriteria">The object containing the criteria to match.</param>
        /// <param name="excludeParent">Exclude the parent element from the search.</param>
        /// <returns>A collection of elements that satisfy the specified conditions. 
        /// If there are no matches, an empty collection is returned.</returns>
        public static List<AutomationElement> FindAll(AutomationElement parentElement, IAutomationElementSearchCriteria searchCriteria, bool excludeParent)
        {
            return Find(parentElement, searchCriteria, false, excludeParent);
        }

        /// <summary>
        /// Returns all descendant elements that match the specified search conditions. 
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="searchCriteria">The object containing the criteria to match.</param>
        /// <param name="milliseconds">The maximum time in milliseconds to allow for attempting to find a successful match.</param>
        /// <returns>
        /// A collection of elements that statisfy the specified search conditions. 
        /// If there are no matches found within the allotted time, an empty collection is returned.
        /// </returns>
        public static List<AutomationElement> FindAll(AutomationElement parentElement, IAutomationElementSearchCriteria searchCriteria, int milliseconds)
        {
            return FindAll(parentElement, searchCriteria, milliseconds, true); 
        }

        /// <summary>
        /// Returns all descendant elements that match the specified search conditions. 
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="searchCriteria">The object containing the criteria to match.</param>
        /// <param name="milliseconds">The maximum time in milliseconds to allow for attempting to find a successful match.</param>
        /// <param name="excludeParent">Exclude parent element from the search.</param>
        /// <returns>
        /// A collection of elements that statisfy the specified search conditions. 
        /// If there are no matches found within the allotted time, an empty collection is returned.
        /// </returns>
        public static List<AutomationElement> FindAll(AutomationElement parentElement, IAutomationElementSearchCriteria searchCriteria, int milliseconds, bool excludeParent)
        {
            var timer = new AutomationTimer(milliseconds);
            timer.Start();
            while (!timer.HasTimedOut())
            {
                var foundElements = Find(parentElement, searchCriteria, false, excludeParent);
                if (foundElements.Count > 0)
                {
                    return foundElements;
                }
            }
            return new List<AutomationElement>();
        }

        /// <summary>
        /// Returns all children elements of a specified <see cref="ControlType"/>.
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="controlType">The <see cref="ControlType"/> of the children.</param>
        /// <returns>A collection of children elements that match the specified control type. 
        /// If there are no matches, an empty collection is returned.
        /// </returns>
        public static List<AutomationElement> FindAllChildrenOfType(AutomationElement parentElement, ControlType controlType)
        {
            var children = new List<AutomationElement>(); 
            TreeWalker treeWalker = TreeWalker.RawViewWalker;
            AutomationElement child = treeWalker.GetFirstChild(parentElement);
            while (child != null)
            {
                if (child.Current.ControlType == controlType)
                {
                    children.Add(child); 
                }
                child = treeWalker.GetNextSibling(child); 
            }
            return children; 
        }

        /// <summary>
        /// Returns all descendants elements of a specified <see cref="ControlType"/>.
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="controlType">The <see cref="ControlType"/> of the descedants.</param>
        /// <returns>A collection of descendants elements that match the specified control type. 
        /// If there are no matches, an empty collection is returned.
        /// </returns>
        public static List<AutomationElement> FindAllDescendantsOfType(AutomationElement parentElement, ControlType controlType)
        {
            var searchCriteria = new AutomationElementSearchCriteria();
            var searchCondition = new AutomationElementExactSearchCondition(AutomationElement.ControlTypeProperty, controlType);
            searchCriteria.Add(searchCondition);
            return FindAll(parentElement, searchCriteria, true); 
        }

        /// <summary>
        /// Finds an element(s) matching the search criteria. 
        /// </summary>
        /// <param name="parentElement">The parent element to search from.</param>
        /// <param name="searchCriteria">The search criteria.</param>
        /// <param name="stopOnFirstMatch">True to indicate to return once a match is found; otherwise continue to search until all matches found.</param>
        /// <param name="excludeParent">Exclude the parent element from search.</param>
        /// <returns>A collection of elements that satisfies the specified search conditions. If there are no matches, an empty collection is returned.</returns>
        /// <exception cref="ArgumentNullException">Thrown if <see cref="parentElement"/> or <see cref="searchCriteria"/> is null.</exception>
        private static List<AutomationElement> Find(AutomationElement parentElement, IAutomationElementSearchCriteria searchCriteria, bool stopOnFirstMatch, bool excludeParent)
        {
            ValidationUtility.ValidateArgumentNotNull(parentElement, "parentElement");
            ValidationUtility.ValidateArgumentNotNull(searchCriteria, "searchCriteria");

            TreeWalker treeWalker = AutomationConfiguration.Current.TreeWalker;

            var foundElements = new List<AutomationElement>();
            var elements = new Queue<AutomationElement>();
            
            if (excludeParent)
            {
                AutomationElement currentChild = treeWalker.GetFirstChild(parentElement);
                while (currentChild != null)
                {
                    elements.Enqueue(currentChild);
                    currentChild = treeWalker.GetNextSibling(currentChild);
                }
            }
            else
            {
                elements.Enqueue(parentElement);
            }
            
            while (elements.Count != 0)
            {
                try
                {
                    AutomationElement currentElement = elements.Dequeue();
                    
                    if (searchCriteria.Matches(currentElement))
                    {
                        foundElements.Add(currentElement);

                        if (searchCriteria.Index >= 0 && searchCriteria.Index == foundElements.Count - 1)
                        {
                            return new List<AutomationElement> { currentElement };
                        }

                        if (stopOnFirstMatch)
                        {
                            return foundElements;
                        }
                    }

                    AutomationElement currentChild = treeWalker.GetFirstChild(currentElement);
                    while (currentChild != null)
                    {
                        elements.Enqueue(currentChild);
                        currentChild = treeWalker.GetNextSibling(currentChild);
                    }
                }
                catch (ElementNotAvailableException) 
                {
                    continue;
                } 
            }
            return foundElements;
        }
    }
}
